from langchain_core.tools import tool
from pydantic import BaseModel, Field
from typing import List, Dict, Any
from langchain_core.runnables import RunnableConfig
from agent.tools.query_indicator_definition import query_indicator_definition
from agent.tools.sql_query import sql_query

class DimensionDrilldownInput(BaseModel):
    indicator: int = Field(description="指标ID")
    dimension: str = Field(description="下钻的维度")
    time_range: List[str] = Field(description="时间范围，格式 ['start_date', 'end_date']")

# 维度下钻工具
@tool("dimension_drilldown", args_schema=DimensionDrilldownInput, return_direct=True)
async def dimension_drilldown(indicator: int, dimension: str, time_range: List[str]) -> Dict[str, Any]:
    """根据维度拆解指标，使用 `sql_query` 工具查询数据"""
    try:
        # 获取指标定义
        indicator_definition = await query_indicator_definition(indicator)
        computation = indicator_definition["computation"]
        data_source = indicator_definition["data_source"]
        filters = indicator_definition["filters"]
        breakdowns = indicator_definition["breakdowns"]
        
        # 构建时间范围过滤条件
        start_date, end_date = time_range
        filters = [{"field": "transaction_date", "condition": f"BETWEEN '{start_date}' AND '{end_date}'"}]

        # 构建 GROUP BY 字段
        group_by = [dimension] if dimension in breakdowns else []

        # 使用 sql_query 工具查询数据
        drilldown_result = await sql_query(computation=computation, data_source=data_source, filters=filters, group_by=group_by)
        
        # 返回查询结果
        return {dimension: drilldown_result}

    except Exception as e:
        print(f"Error in dimension drilldown: {e}")
        return {}
